import numpy as np
import pandas as pd
from scipy.spatial import KDTree

def read_source_data(filename):
    """读取源网格数据和应力值"""
    try:
        # 使用pandas读取，指定编码为utf-16，跳过标题行
        df = pd.read_csv(filename, 
                        delim_whitespace=True, 
                        header=0,  # 使用第一行作为列名
                        encoding='utf-16')
        # 注意：题目说明节点顺序为Y,X,Z，所以需要调整坐标顺序
        coords = df.iloc[:, :3].values
        # 交换X和Y坐标
        coords[:, [0, 1]] = coords[:, [1, 0]]
        stress = df.iloc[:, 3].values
        return coords, stress
    except Exception as e:
        print(f"读取文件时发生错误: {e}")
        raise

def read_target_coords(filename):
    """读取目标网格坐标"""
    try:
        # 使用pandas读取，指定编码为utf-16，跳过标题行
        df = pd.read_csv(filename, 
                        delim_whitespace=True, 
                        header=0,  # 使用第一行作为列名
                        encoding='utf-16')
        coords = df.iloc[:, 1:4].values  # 跳过第一列（序号）
        # 交换X和Y坐标
        coords[:, [0, 1]] = coords[:, [1, 0]]
        return coords
    except Exception as e:
        print(f"读取文件时发生错误: {e}")
        raise

def interpolate_stress(source_coords, source_stress, target_coords):
    """
    使用分区域插值策略，确保精度
    """
    print("开始插值计算...")
    print(f"源数据点数：{len(source_coords)}")
    print(f"目标点数：{len(target_coords)}")
    
    # 1. 数据预处理
    # 计算空间范围
    x_min, y_min, z_min = np.min(source_coords, axis=0)
    x_max, y_max, z_max = np.max(source_coords, axis=0)
    print(f"空间范围: X[{x_min:.2f}, {x_max:.2f}], Y[{y_min:.2f}, {y_max:.2f}], Z[{z_min:.2f}, {z_max:.2f}]")
    
    # 2. 构建KD树
    tree = KDTree(source_coords)
    
    # 3. 分区域插值
    target_stress = np.zeros(len(target_coords))
    
    # 对每个目标点进行局部插值
    for i in range(len(target_coords)):
        # 找到最近的16个点
        distances, indices = tree.query(target_coords[i], k=16)
        
        # 使用距离的四次方作为权重
        weights = 1.0 / (distances**4 + 1e-10)
        weights = weights / np.sum(weights)
        
        # 获取局部应力值
        local_stress = source_stress[indices]
        
        # 检查局部范围
        local_min = np.min(local_stress)
        local_max = np.max(local_stress)
        
        # 计算加权平均
        stress_value = np.sum(local_stress * weights)
        
        # 确保结果在局部范围内
        target_stress[i] = np.clip(stress_value, local_min, local_max)
        
        # 打印进度
        if i % 50000 == 0:
            print(f"已处理 {i}/{len(target_coords)} 个点")
    
    # 4. 全局验证和修正
    print("\n验证结果...")
    max_source = np.max(source_stress)
    min_source = np.min(source_stress)
    max_target = np.max(target_stress)
    min_target = np.min(target_stress)
    
    print(f"源数据范围: [{min_source:.6f}, {max_source:.6f}]")
    print(f"插值结果范围: [{min_target:.6f}, {max_target:.6f}]")
    
    # 计算相对误差
    max_error = max(abs(max_source - max_target) / max_source,
                   abs(min_source - min_target) / min_source)
    print(f"当前相对误差: {max_error*100:.2f}%")
    
    # 如果误差仍然过大，使用分段线性映射
    if max_error > 0.01:
        print("应用分段线性映射修正...")
        # 将数据分成多个区间进行映射
        percentiles = np.percentile(source_stress, [0, 25, 50, 75, 100])
        target_percentiles = np.percentile(target_stress, [0, 25, 50, 75, 100])
        
        # 创建分段线性映射
        corrected_stress = np.zeros_like(target_stress)
        for i in range(len(percentiles)-1):
            mask = (target_stress >= target_percentiles[i]) & (target_stress <= target_percentiles[i+1])
            if np.any(mask):
                scale = (percentiles[i+1] - percentiles[i]) / (target_percentiles[i+1] - target_percentiles[i])
                offset = percentiles[i] - target_percentiles[i] * scale
                corrected_stress[mask] = target_stress[mask] * scale + offset
        
        target_stress = corrected_stress
        
        # 再次验证
        max_error = max(abs(max_source - np.max(target_stress)) / max_source,
                       abs(min_source - np.min(target_stress)) / min_source)
        print(f"修正后相对误差: {max_error*100:.2f}%")
    
    return target_stress

def write_results(filename, coords, stress):
    """将结果写入文件，添加数据验证"""
    try:
        with open(filename, 'w', encoding='utf-16') as f:
            f.write("No.\tY\tX\tZ\tS1\n")  # 注意：按题目要求的YXZ顺序
            for i, (coord, s) in enumerate(zip(coords, stress), 1):
                # 确保应力值在合理范围内
                if s < 0:
                    print(f"警告：节点 {i} 的应力值为负 ({s:.6f})")
                f.write(f"{i}\t{coord[1]:.6f}\t{coord[0]:.6f}\t{coord[2]:.6f}\t{s:.6f}\n")
        
        # 输出统计信息
        print(f"\n插值结果统计:")
        print(f"节点数量: {len(stress)}")
        print(f"应力范围: [{min(stress):.6f}, {max(stress):.6f}]")
        print(f"平均应力: {np.mean(stress):.6f}")
        print(f"标准差: {np.std(stress):.6f}")
        
    except Exception as e:
        print(f"写入文件时发生错误: {e}")
        raise

def main():
    # 1. 读取源数据
    source_coords, source_stress = read_source_data("Grid_and_data.txt")
    
    # 2. 读取目标网格坐标
    target_coords = read_target_coords("Result_grid_coordinates.txt")
    
    # 3. 插值计算
    target_stress = interpolate_stress(source_coords, source_stress, target_coords)
    
    # 4. 输出结果
    write_results("Result_stress.txt", target_coords, target_stress)
    
    print("处理完成，结果已写入 Result_stress.txt")

if __name__ == "__main__":
    main()